# 给定一个整数数组，返回所有数对之间的第 k 个最小距离。一对 (A, B) 的距离被定义为 A 和 B 之间的绝对差值。
#
#  示例 1:
# 输入：
# nums = [1,3,1]
# k = 1
# 输出：0
# 解释：
# 所有数对如下：
# (1,3) -> 2
# (1,1) -> 0
# (3,1) -> 2
# 因此第 1 个最小距离的数对是 (1,1)，它们之间的距离为 0。
import heapq
from typing import List


class Solution:
    def smallestDistancePair2(self, nums: List[int], k: int) -> int:
        """
        解法二：二分查找 + 双指针
        由于第 k 小的距离一定在 [0,  high = max(nums) - min(nums)] 中，我们在这个区间上进行二分。
        对于当前二分的位置 mid，统计距离小于等于 mid 的距离对数量用count表示(双指针实现)。
        如果count > k 则从(low, mid - 1)继续二分查找
        如果count < k 则从(mid + 1, high)继续二分查找
        如果count = k 则返回 mid
        :param nums:
        :param k:
        :return:
        """
        sortedNums = sorted(nums)
        numsLen = len(nums)
        low, high = 0, sortedNums[-1] - sortedNums[0]
        while low < high:
            mid = low + ((high - low) >> 1)
            count, left, right = 0, 0, 0
            # 双指针统计距离小于等于 mid 的距离对数量
            while right < numsLen:
                while sortedNums[right] - sortedNums[left] > mid:
                    left += 1
                count += right - left
                right += 1

            # if count > k:  # 如果符合距离对小于等于 mid 的数对个数太多了，那么就从左边查找
            #     high = mid - 1
            # elif count > k:  # 如果符合距离对小于等于 mid 的数对个数太少了，那么就从右边查找
            #     low = mid + 1
            # else:
            #     return low
            if count >= k:  # 如果符合距离对小于等于 mid 的数对个数太多了，那么就从左边查找
                high = mid
            else:  # 如果符合距离对小于等于 mid 的数对个数太少了，那么就从右边查找
                low = mid + 1
        return low

    def smallestDistancePair1(self, nums: List[int], k: int) -> int:
        """
        解法一：最小堆(超时)
        先对数组进行排序，则：(i, i+1)的距离一定是小于等于(i, i + 2), (i, i + 3), (i, i + n)
        因为 (i, j + 1) 的距离不会小于 (i, j)，因此只要 (i, j) 还在堆中，就没有必要把 (i, j + 1) 放入堆中。
        然后将所有 (i, i + 1) 放入堆中，
        每次取出堆中的最小元素 (i, j) 时（元素的大小为 nums[j] - nums[i]，即距离），再把 (i, j + 1) 放入堆中。
        直到取出 k 个元素，就得到了第 k 小的距离。
        :param nums:
        :param k:
        :return:
        """
        sortedNums = sorted(nums)
        numsLen = len(nums)
        heap = []
        for i in range(numsLen - 1):
            heapq.heappush(heap, (sortedNums[i + 1] - sortedNums[i], i, i + 1))  # 最小的距离对一定在所有的(i, i + 1)对中
        dist, index1, index2 = 0, 0, 0
        while k:  # 直到取出 k 个元素，就得到了第 k 小的距离。
            dist, index1, index2 = heapq.heappop(heap)
            if index2 < numsLen - 1:  # 每次取出堆中的最小元素 (i, j) 时（元素的大小为 nums[j] - nums[i]，即距离），再把 (i, j + 1) 放入堆中。
                heapq.heappush(heap, (sortedNums[index2 + 1] - sortedNums[index1], index1, index2 + 1))
            k -= 1
        return sortedNums[index2] - sortedNums[index1]

    def smallestDistancePair(self, nums: List[int], k: int) -> int:
        return self.smallestDistancePair2(nums, k)


if __name__ == "__main__":
    # nums = [1, 3, 1]
    # k = 1
    nums = [1, 6, 1]
    k = 3
    print(Solution().smallestDistancePair(nums, k))
